package com.ai.chat.service;

import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import com.ai.chat.adapter.FullChatAdapter;
import com.ai.chat.domain.base.request.BaseChatRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;

@Service
public class McpService extends FullChatAdapter {
    private final ChatClient defaultChatClient;

    private final ConcurrentHashMap<String, Disposable> streamDisposableMap = new ConcurrentHashMap<>();

    @Autowired
    public McpService(ChatClient.Builder chatClientBuilder, ToolCallbackProvider toolCallbackProvider) {
        for (FunctionCallback toolCallback : toolCallbackProvider.getToolCallbacks()) {
            System.out.println(toolCallback.getName());
        }
        this.defaultChatClient = chatClientBuilder
                .defaultTools(toolCallbackProvider)
                .build();
    }

    @Override
    public String chat(BaseChatRequest request) {
        return defaultChatClient.prompt(request.getMessage()).call().content();
    }

    @Override
    public Flux<String> streamChat(BaseChatRequest request, HttpServletResponse response) {
        String streamId = UUID.randomUUID().toString().replaceAll("-", "");
        response.setHeader("X-Stream-Id", streamId);

        Flux<String> answer = defaultChatClient.prompt().user(request.getMessage()).stream().content();

        return answer.doOnSubscribe(sub -> {
                    // 将Subscription转换为Disposable
                    Disposable disposable = sub::cancel;
                    streamDisposableMap.put(streamId, disposable);
                })
                .doOnCancel(() -> disposeStream(streamId))
                .doFinally(signal -> streamDisposableMap.remove(streamId));
    }

    @Override
    public void disposeStream(String streamId) {
        Disposable disposable = streamDisposableMap.get(streamId);
        if (disposable != null && !disposable.isDisposed()) {
            disposable.dispose();
            streamDisposableMap.remove(streamId);
        }
    }
}
